#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <sensor_msgs/msg/camera_info.hpp>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <cstring>
#include <vector>

class CameraNetworkBridge : public rclcpp::Node {
public:
    CameraNetworkBridge() : Node("camera_network_bridge") {
        this->declare_parameter<std::string>("target_ip", "192.168.3.238");
        this->declare_parameter<int>("target_port", 50010);
        target_ip_ = this->get_parameter("target_ip").as_string();
        target_port_ = this->get_parameter("target_port").as_int();

        // 主TCP socket (50010)
        sock_fd_ = socket(AF_INET, SOCK_STREAM, 0);
        if (sock_fd_ < 0) {
            RCLCPP_ERROR(this->get_logger(), "TCP socket create failed");
            return;
        }
        memset(&server_addr_, 0, sizeof(server_addr_));
        server_addr_.sin_family = AF_INET;
        server_addr_.sin_port = htons(target_port_);
        server_addr_.sin_addr.s_addr = inet_addr(target_ip_.c_str());
        if (connect(sock_fd_, (struct sockaddr*)&server_addr_, sizeof(server_addr_)) < 0) {
            RCLCPP_ERROR(this->get_logger(), "TCP connect failed to %s:%d", target_ip_.c_str(), target_port_);
            close(sock_fd_);
            sock_fd_ = -1;
        } else {
            RCLCPP_INFO(this->get_logger(), "TCP connected to %s:%d", target_ip_.c_str(), target_port_);
        }

        // 新增RGB专用TCP socket (50011)
        sock_fd_rgb_ = socket(AF_INET, SOCK_STREAM, 0);
        if (sock_fd_rgb_ < 0) {
            RCLCPP_ERROR(this->get_logger(), "TCP socket (RGB) create failed");
            sock_fd_rgb_ = -1;
        } else {
            memset(&rgb_addr_, 0, sizeof(rgb_addr_));
            rgb_addr_.sin_family = AF_INET;
            rgb_addr_.sin_port = htons(50011);
            rgb_addr_.sin_addr.s_addr = inet_addr(target_ip_.c_str());
            if (connect(sock_fd_rgb_, (struct sockaddr*)&rgb_addr_, sizeof(rgb_addr_)) < 0) {
                RCLCPP_ERROR(this->get_logger(), "TCP connect (RGB) failed to %s:%d", target_ip_.c_str(), 50011);
                close(sock_fd_rgb_);
                sock_fd_rgb_ = -1;
            } else {
                RCLCPP_INFO(this->get_logger(), "TCP (RGB) connected to %s:%d", target_ip_.c_str(), 50011);
            }
        }

        // 订阅相机话题
        color_img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/camera/color/image_raw", 10,
            [this](sensor_msgs::msg::Image::SharedPtr msg) { image_callback(msg, 0); });

        color_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
            "/camera/color/camera_info", 10,
            [this](sensor_msgs::msg::CameraInfo::SharedPtr msg) { info_callback(msg, 0); });

        depth_img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/camera/depth/image_raw", 10,
            [this](sensor_msgs::msg::Image::SharedPtr msg) { image_callback(msg, 1); });

        depth_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
            "/camera/depth/camera_info", 10,
            [this](sensor_msgs::msg::CameraInfo::SharedPtr msg) { info_callback(msg, 1); });
    }
    ~CameraNetworkBridge() {
        if (sock_fd_ >= 0) close(sock_fd_);
        if (sock_fd_rgb_ >= 0) close(sock_fd_rgb_);
    }
private:
    // type: 0=color, 1=depth
    void image_callback(const sensor_msgs::msg::Image::SharedPtr msg, int type) {
        std::vector<uint8_t> buffer;
        buffer.push_back(type);
        int64_t stamp = msg->header.stamp.sec * 1000000000LL + msg->header.stamp.nanosec;
        buffer.insert(buffer.end(), (uint8_t*)&stamp, (uint8_t*)&stamp + sizeof(stamp));
        buffer.insert(buffer.end(), (uint8_t*)&msg->height, (uint8_t*)&msg->height + sizeof(msg->height));
        buffer.insert(buffer.end(), (uint8_t*)&msg->width, (uint8_t*)&msg->width + sizeof(msg->width));
        std::string enc = msg->encoding;
        buffer.insert(buffer.end(), enc.begin(), enc.end());
        buffer.push_back(0);
        buffer.insert(buffer.end(), msg->data.begin(), msg->data.end());
        // 主socket发送
        if (sock_fd_ >= 0) {
            ssize_t sent = send(sock_fd_, buffer.data(), buffer.size(), 0);
            if (sent > 0) {
                RCLCPP_INFO(this->get_logger(), "[NET] Sent image type %d, size %ld bytes", type, sent);
            } else {
                RCLCPP_WARN(this->get_logger(), "[NET] Failed to send image type %d", type);
            }
        }
        // 如果是RGB图像，额外发到50011端口
        if (type == 0 && sock_fd_rgb_ >= 0) {
            ssize_t sent_rgb = send(sock_fd_rgb_, buffer.data(), buffer.size(), 0);
            if (sent_rgb > 0) {
                RCLCPP_INFO(this->get_logger(), "[NET] Sent RGB image to 50011, size %ld bytes", sent_rgb);
            } else {
                RCLCPP_WARN(this->get_logger(), "[NET] Failed to send RGB image to 50011");
            }
        }
    }
    void info_callback(const sensor_msgs::msg::CameraInfo::SharedPtr msg, int type) {
        std::vector<uint8_t> buffer;
        buffer.push_back(type+10);
        int64_t stamp = msg->header.stamp.sec * 1000000000LL + msg->header.stamp.nanosec;
        buffer.insert(buffer.end(), (uint8_t*)&stamp, (uint8_t*)&stamp + sizeof(stamp));
        buffer.insert(buffer.end(), (uint8_t*)&msg->height, (uint8_t*)&msg->height + sizeof(msg->height));
        buffer.insert(buffer.end(), (uint8_t*)&msg->width, (uint8_t*)&msg->width + sizeof(msg->width));
        for (int i = 0; i < 9; ++i) {
            buffer.insert(buffer.end(), (uint8_t*)&msg->k[i], (uint8_t*)&msg->k[i] + sizeof(double));
        }
        if (sock_fd_ >= 0) {
            ssize_t sent = send(sock_fd_, buffer.data(), buffer.size(), 0);
            if (sent > 0) {
                RCLCPP_INFO(this->get_logger(), "[NET] Sent camera_info type %d, size %ld bytes", type+10, sent);
            } else {
                RCLCPP_WARN(this->get_logger(), "[NET] Failed to send camera_info type %d", type+10);
            }
        }
    }
    int sock_fd_ = -1;
    int sock_fd_rgb_ = -1;
    struct sockaddr_in server_addr_;
    struct sockaddr_in rgb_addr_;
    std::string target_ip_;
    int target_port_;
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr color_img_sub_;
    rclcpp::Subscription<sensor_msgs::msg::CameraInfo>::SharedPtr color_info_sub_;
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr depth_img_sub_;
    rclcpp::Subscription<sensor_msgs::msg::CameraInfo>::SharedPtr depth_info_sub_;
};

int main(int argc, char** argv) {
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<CameraNetworkBridge>());
    rclcpp::shutdown();
    return 0;
}
